# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/29 15:28 
@Author  : ZhangShenao 
@File    : 5.使用自定义检索器.py 
@Desc    : 使用SimpleRetriever自定义检索器
"""
from langchain_core.documents import Document

from simple_retriever import SimpleRetriever

# 构造文档列表
docs = [
    Document(page_content="笨笨是一只很喜欢睡觉的猫咪", metadata={"page": 1}),
    Document(page_content="我喜欢在夜晚听音乐，这让我感到放松。", metadata={"page": 2}),
    Document(page_content="猫咪在窗台上打盹，看起来非常可爱。", metadata={"page": 3}),
    Document(page_content="学习新技能是每个人都应该追求的目标。", metadata={"page": 4}),
    Document(page_content="我最喜欢的食物是意大利面，尤其是番茄酱的那种。", metadata={"page": 5}),
    Document(page_content="昨晚我做了一个奇怪的梦，梦见自己在太空飞行。", metadata={"page": 6}),
    Document(page_content="我的手机突然关机了，让我有些焦虑。", metadata={"page": 7}),
    Document(page_content="阅读是我每天都会做的事情，我觉得很充实。", metadata={"page": 8}),
    Document(page_content="他们一起计划了一次周末的野餐，希望天气能好。", metadata={"page": 9}),
    Document(page_content="我的狗喜欢追逐球，看起来非常开心。", metadata={"page": 10}),
]

# 创建自定义检索器
retriever = SimpleRetriever(documents=docs, top_k=3)

# 执行自定义检索器
# BaseRetriever也是Runnable可运行组件,可以通过LCEL表达式直接运行
result = retriever.invoke('猫')
print(result)
